package org.mybatis.generator.codegen.spring.jdbc.repository.elements;

import static org.mybatis.generator.internal.util.JavaBeansUtil.getSetterMethodName;

import java.util.List;

import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.dom.java.FullyQualifiedJavaType;
import org.mybatis.generator.api.dom.java.JavaVisibility;
import org.mybatis.generator.api.dom.java.Method;
import org.mybatis.generator.api.dom.java.Parameter;
import org.mybatis.generator.api.dom.java.TopLevelClass;
import org.mybatis.generator.codegen.spring.jdbc.repository.AbstractRepGenerator;

public class SelectByPrimaryKeyMethodGenerator extends AbstractRepGenerator {

	@Override
	public void addImplementationElements(TopLevelClass topLevelClass) {
		// 添加依赖类
		addImportedType(topLevelClass);
		Boolean isadd = genFindByKey(topLevelClass);
		if (isadd) {
			//
			addSelectByPrimaryKeysMethodElements(topLevelClass);
		}
	}

	private Boolean genFindByKey(TopLevelClass topLevelClass) {
		// 获得主键信息
		List<IntrospectedColumn> pks = introspectedTable.getPrimaryKeyColumns();
		// IntrospectedColumn r =
		// introspectedTable.getPrimaryKeyColumns().get(0);
		StringBuilder sb = new StringBuilder();
		// 设置方法
		Method method = new Method();
		method.addException(new FullyQualifiedJavaType("org.springframework.dao.DataAccessException"));
		method.setVisibility(JavaVisibility.PUBLIC);
		FullyQualifiedJavaType returnType = introspectedTable.getRules().calculateAllFieldsClass();
		method.setReturnType(returnType);
//		method.setReturnType(new FullyQualifiedJavaType(introspectedTable.getFullyQualifiedTable()
//		        .getDomainObjectName()));
		topLevelClass.addImportedType(returnType);
		method.setName("findOneById");

		// 方法内部代码
		method.addBodyLine(returnType.getShortName() + " " + "entity;");
		method.addBodyLine("try {");
		method.addBodyLine("Map<String, Object> params = new HashMap<String, Object>();");
		for (IntrospectedColumn introspectedColumn : pks) {
			String column = toCaseName(introspectedColumn.getActualColumnName());
			// 添加参数
			method.addParameter(new Parameter(introspectedColumn.getFullyQualifiedJavaType(), column));
			StringBuffer sb_ = new StringBuffer("params.put(\"");
			sb_.append(column);
			sb_.append("\",");
			sb_.append(column);
			sb_.append(");");
			method.addBodyLine(sb_.toString());
		}
		sb.setLength(0);
		sb.append("entity = this.jdbcTemplate.queryForObject(\"");
		sb.append(getSqlSelect(introspectedTable));
		sb.append(" WHERE ");
		for (int i = 0; i < pks.size(); i++) {
			IntrospectedColumn introspectedColumn = pks.get(i);
			if (i != 0) {
				sb.append(" and ");
			}
			String column = toCaseName(introspectedColumn.getActualColumnName());
			sb.append(column);
			sb.append(" = :");
			sb.append(column);
		}
		sb.append("\",params, BeanPropertyRowMapper.newInstance(");
		sb.append(returnType.getShortName());
		sb.append(".class));");
		method.addBodyLine(sb.toString());

		method.addBodyLine("} catch (EmptyResultDataAccessException ex) {");
		method.addBodyLine("throw ex;\n\r}");
		method.addBodyLine("return entity;");
		context.getCommentGenerator().addGeneralMethodComment(method, introspectedTable);

		if (context.getPlugins().clientSelectByPrimaryKeyMethodGenerated(method, topLevelClass, introspectedTable)) {
			topLevelClass.addMethod(method);
			return true;
		}
		return false;
	}

	/**
	 * @param topLevelClass
	 */
	private void addImportedType(TopLevelClass topLevelClass) {
		// 添加依赖类
		topLevelClass.addImportedType(new FullyQualifiedJavaType("java.util.Map"));
		topLevelClass.addImportedType(new FullyQualifiedJavaType("java.util.HashMap"));
		topLevelClass.addImportedType(new FullyQualifiedJavaType("java.util.List"));

		topLevelClass.addImportedType(new FullyQualifiedJavaType("java.sql.SQLException"));
		topLevelClass.addImportedType(new FullyQualifiedJavaType("java.sql.ResultSet"));

		topLevelClass
		        .addImportedType(new FullyQualifiedJavaType("org.springframework.jdbc.core.BeanPropertyRowMapper"));
		topLevelClass.addImportedType(new FullyQualifiedJavaType("org.springframework.jdbc.core.RowMapper"));
		topLevelClass.addImportedType(new FullyQualifiedJavaType("org.springframework.dao.DataAccessException"));
		topLevelClass.addImportedType(new FullyQualifiedJavaType(
		        "org.springframework.dao.EmptyResultDataAccessException"));
		FullyQualifiedJavaType modelClass = new FullyQualifiedJavaType(introspectedTable.getBaseRecordType());
		topLevelClass.addImportedType(modelClass);
		topLevelClass.addImportedType(new FullyQualifiedJavaType(
		        "org.springframework.jdbc.core.namedparam.EmptySqlParameterSource"));
		topLevelClass.addImportedType(new FullyQualifiedJavaType(
		        "org.springframework.jdbc.core.namedparam.MapSqlParameterSource"));
		topLevelClass.addImportedType(new FullyQualifiedJavaType(
		        "org.springframework.jdbc.core.namedparam.SqlParameterSource"));

	}

	private void addSelectByPrimaryKeysMethodElements(TopLevelClass topLevelClass) {
		List<IntrospectedColumn> pks = introspectedTable.getPrimaryKeyColumns();
		if (pks.size() > 1) {
			// 如果是复合主键，不生成多主键查询方法
			return;
		}

		//String domainObjectName = introspectedTable.getFullyQualifiedTable().getDomainObjectName();
		FullyQualifiedJavaType returnType_ = introspectedTable.getRules().calculateAllFieldsClass();
		String domainObjectName = returnType_.getShortName();
		// String primaryKeyTypeName = introspectedTable.getPrimaryKeyType();
		// 获得主键信息
		IntrospectedColumn r = pks.get(0);
		String primaryKeyTypeName = r.getFullyQualifiedJavaType().getShortName();

		FullyQualifiedJavaType returnType = new FullyQualifiedJavaType("List<" + domainObjectName + ">");
		FullyQualifiedJavaType parameterType = new FullyQualifiedJavaType("List<" + primaryKeyTypeName + ">");

		// 设置方法
		Method method = new Method();
		method.addException(new FullyQualifiedJavaType("org.springframework.dao.DataAccessException"));
		method.setVisibility(JavaVisibility.PUBLIC);
		method.setReturnType(returnType);
		method.setName("findByIds");
		method.addParameter(new Parameter(parameterType, "ids"));
		// 方法内部代码
		StringBuilder sb = new StringBuilder();

		method.addBodyLine(returnType.toString() + " " + "list;");
		method.addBodyLine("SqlParameterSource params = null;");
		method.addBodyLine("String  sql = null;");
		method.addBodyLine("if(ids== null || ids.isEmpty()){");
		method.addBodyLine("params = EmptySqlParameterSource.INSTANCE;");
		method.addBodyLine("sql = \"" + getSqlSelect(introspectedTable) + "\";");
		method.addBodyLine("}else{");
		method.addBodyLine("params = new MapSqlParameterSource(\"ids\",ids);");
		sb.setLength(0);
		sb.append("sql = \"");
		sb.append(getSqlSelect(introspectedTable));
		sb.append(" WHERE ");
		sb.append(r.getActualColumnName());
		sb.append(" in (:ids) \";");
		method.addBodyLine(sb.toString());
		method.addBodyLine("}");

		sb.setLength(0);
		sb.append("list = this.jdbcTemplate.query(sql,params,  new RowMapper<");
		sb.append(domainObjectName);
		sb.append(">() {");
		method.addBodyLine(sb.toString());

		sb.setLength(0);
		sb.append("public ");
		sb.append(domainObjectName);
		sb.append(" mapRow(ResultSet rs, int rowNum) throws SQLException { ");
		method.addBodyLine(sb.toString());

		sb.setLength(0);
		sb.append(domainObjectName);
		sb.append(" row = new ");
		sb.append(domainObjectName);
		sb.append("();");
		method.addBodyLine(sb.toString());

		List<IntrospectedColumn> iterable = introspectedTable.getAllColumns();
		for (int i = 0, size = iterable.size(); i < size; i++) {
			IntrospectedColumn introspectedColumn = iterable.get(i);
			sb.setLength(0);
			sb.append(" row.");
			sb.append(getSetterMethodName(introspectedColumn.getJavaProperty()));
			sb.append("(rs.get");
			sb.append(getRsType(introspectedColumn.getFullyQualifiedJavaType().getShortName()));
			sb.append("(\"");
			sb.append(introspectedColumn.getJavaProperty());
			sb.append("\"));");
			method.addBodyLine(sb.toString());
		}

		method.addBodyLine("return row;");
		method.addBodyLine("}");
		method.addBodyLine("});");
		method.addBodyLine("return list;");
		context.getCommentGenerator().addGeneralMethodComment(method, introspectedTable);
		topLevelClass.addMethod(method);
	}

	private String getRsType(String shortName) {
		if (shortName.equals("Integer")) {
			return "Int";
		} else if ("byte[]".equals(shortName)) {
			return "Bytes";
		}
		return shortName;
	}

	private String getSqlSelect(IntrospectedTable introspectedTable) {
		List<IntrospectedColumn> iterable = introspectedTable.getAllColumns();
		StringBuilder sb = new StringBuilder("SELECT ");
		for (int i = 0, size = iterable.size(); i < size; i++) {
			IntrospectedColumn element = iterable.get(i);
			sb.append(element.getActualColumnName().toLowerCase());
			sb.append(" AS ");
			sb.append(element.getJavaProperty());
			if (i + 1 < size) {
				sb.append(", ");
			}
		}
		sb.append(" FROM ");
		sb.append(introspectedTable.getFullyQualifiedTable().getIntrospectedTableName());
		return sb.toString();
	}
}
